import sqlite3
import os
from typing import List, Dict, Any, Optional
from pathlib import Path
from ..core.logger import get_logger

logger = get_logger('agent.db_tool')

class DatabaseTool:
    """数据库访问工具类"""
    
    def __init__(self, db_path: str = None):
        """
        初始化数据库工具
        
        Args:
            db_path: 数据库文件路径
        """
        # 默认数据库路径
        if db_path is None:
            self.db_path = Path(__file__).parent.parent.parent.parent / "dataset" / "sqlite"
        else:
            self.db_path = Path(db_path)
        
        # 确保目录存在
        self.db_path.mkdir(parents=True, exist_ok=True)
        
        logger.info(f"Database tool initialized with path: {self.db_path}")
    
    def list_databases(self) -> List[str]:
        """
        列出所有可用的数据库文件
        
        Returns:
            数据库文件名列表
        """
        try:
            db_files = [f for f in os.listdir(self.db_path) if f.endswith('.db') or f.endswith('.sqlite')]
            return db_files
        except Exception as e:
            logger.error(f"Error listing databases: {e}")
            return []
    
    def connect_database(self, db_name: str) -> Optional[sqlite3.Connection]:
        """
        连接到指定数据库
        
        Args:
            db_name: 数据库文件名
            
        Returns:
            数据库连接对象或None
        """
        try:
            db_file = self.db_path / db_name
            if not db_file.exists():
                logger.error(f"Database file not found: {db_file}")
                return None
            
            conn = sqlite3.connect(db_file)
            conn.row_factory = sqlite3.Row  # 使结果可以通过列名访问
            return conn
        except Exception as e:
            logger.error(f"Error connecting to database {db_name}: {e}")
            return None
    
    def execute_query(self, db_name: str, query: str, params: tuple = ()) -> List[Dict[str, Any]]:
        """
        执行SQL查询
        
        Args:
            db_name: 数据库文件名
            query: SQL查询语句
            params: 查询参数
            
        Returns:
            查询结果列表
        """
        conn = self.connect_database(db_name)
        if not conn:
            return []
        
        try:
            cursor = conn.cursor()
            cursor.execute(query, params)
            results = cursor.fetchall()
            
            # 转换为字典列表
            return [dict(row) for row in results]
        except Exception as e:
            logger.error(f"Error executing query on {db_name}: {e}")
            return []
        finally:
            conn.close()
    
    def get_table_names(self, db_name: str) -> List[str]:
        """
        获取数据库中所有表名
        
        Args:
            db_name: 数据库文件名
            
        Returns:
            表名列表
        """
        query = "SELECT name FROM sqlite_master WHERE type='table';"
        results = self.execute_query(db_name, query)
        return [row['name'] for row in results]
    
    def get_table_schema(self, db_name: str, table_name: str) -> List[Dict[str, Any]]:
        """
        获取表的结构信息
        
        Args:
            db_name: 数据库文件名
            table_name: 表名
            
        Returns:
            表结构信息列表
        """
        query = f"PRAGMA table_info('{table_name}');"
        return self.execute_query(db_name, query)

    def create_financial_database(self, db_name: str) -> bool:
        """
        创建金融数据库的默认结构
        
        Args:
            db_name: 数据库文件名
            
        Returns:
            成功创建返回 True，否则 False
        """
        conn = self.connect_database(db_name)
        if not conn:
            return False

        try:
            cursor = conn.cursor()
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS financial_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    date TEXT NOT NULL,
                    stock_symbol TEXT NOT NULL,
                    open REAL,
                    high REAL,
                    low REAL,
                    close REAL,
                    volume INTEGER,
                    UNIQUE(date, stock_symbol)
                )
            ''')
            conn.commit()
            return True
        except Exception as e:
            logger.error(f"Error creating financial database: {e}")
            return False
        finally:
            conn.close()

    def validate_financial_table(self, db_name: str, table_name: str = "financial_data") -> bool:
        """
        验证指定数据库中的表是否符合金融数据的标准结构
        
        Args:
            db_name: 数据库文件名
            table_name: 表名
            
        Returns:
            如果表结构符合返回 True，否则 False
        """
        schema = self.get_table_schema(db_name, table_name)
        required_columns = {
            'date': 'TEXT',
            'stock_symbol': 'TEXT',
            'open': 'REAL',
            'high': 'REAL',
            'low': 'REAL',
            'close': 'REAL',
            'volume': 'INTEGER'
        }

        for col in schema:
            name = col['name']
            type_ = col['type']
            if name in required_columns and required_columns[name] != type_:
                return False

        return True

    def query_financial_data(self, db_name: str, stock_symbol: str, start_date: str = None, end_date: str = None) -> List[Dict[str, Any]]:
        """
        查询金融数据
        
        Args:
            db_name: 数据库文件名
            stock_symbol: 股票代码
            start_date: 开始日期 (格式: YYYY-MM-DD)
            end_date: 结束日期 (格式: YYYY-MM-DD)
            
        Returns:
            查询结果列表
        """
        query = "SELECT * FROM financial_data WHERE stock_symbol = ?"
        params = [stock_symbol]

        if start_date:
            query += " AND date >= ?"
            params.append(start_date)
        if end_date:
            query += " AND date <= ?"
            params.append(end_date)

        return self.execute_query(db_name, query, tuple(params))

# MCP兼容功能
def create_mcp_database_tool(db_path: str = None) -> DatabaseTool:
    """
    创建MCP兼容的数据库工具实例
    
    Args:
        db_path: 数据库路径
        
    Returns:
        DatabaseTool实例
    """
    return DatabaseTool(db_path)

def mcp_list_databases(tool: DatabaseTool) -> List[str]:
    """
    MCP兼容的列出数据库函数
    
    Args:
        tool: DatabaseTool实例
        
    Returns:
        数据库列表
    """
    return tool.list_databases()

def mcp_execute_query(tool: DatabaseTool, db_name: str, query: str) -> List[Dict[str, Any]]:
    """
    MCP兼容的执行查询函数
    
    Args:
        tool: DatabaseTool实例
        db_name: 数据库名
        query: SQL查询语句
        
    Returns:
        查询结果
    """
    return tool.execute_query(db_name, query)

def mcp_get_table_names(tool: DatabaseTool, db_name: str) -> List[str]:
    """
    MCP兼容的获取表名函数
    
    Args:
        tool: DatabaseTool实例
        db_name: 数据库名
        
    Returns:
        表名列表
    """
    return tool.get_table_names(db_name)

def mcp_get_table_schema(tool: DatabaseTool, db_name: str, table_name: str) -> List[Dict[str, Any]]:
    """
    MCP兼容的获取表结构函数
    
    Args:
        tool: DatabaseTool实例
        db_name: 数据库名
        table_name: 表名
        
    Returns:
        表结构信息
    """
    return tool.get_table_schema(db_name, table_name)